# rag_pipeline.py
import requests
from pymilvus import connections, Collection

# 连接 Milvus
def connect_milvus(host="127.0.0.1", port="19530", alias="default"):
    connections.connect(alias=alias, host=host, port=port)
    print(f"✅ 已连接 Milvus ({host}:{port})")

# 生成查询向量
def get_embedding(text, embed_model="nomic-embed-text:latest",
                  ollama_embed_url="http://127.0.0.1:11434/api/embeddings"):
    payload = {"model": embed_model, "prompt": text}
    resp = requests.post(ollama_embed_url, json=payload)
    resp.raise_for_status()
    emb = resp.json().get("embedding")
    if not emb:
        raise RuntimeError("❌ 生成 embedding 失败")
    return emb

# RAG 检索 + 调用大模型生成回答
def rag_query(question, collection_name="doc_chunks",
              embed_model="nomic-embed-text:latest",
              ollama_embed_url="http://127.0.0.1:11434/api/embeddings",
              llama2_url="http://127.0.0.1:3000/api/generate",  # LLaMA2 接口
              llama2_model="llama2:latest",
              top_k=3):
    # 1. 生成向量
    q_emb = get_embedding(question, embed_model, ollama_embed_url)

    # 2. Milvus 检索
    collection = Collection(collection_name)
    collection.load()
    search_params = {"metric_type": "COSINE", "params": {"nprobe": 16}}
    results = collection.search(
        data=[q_emb],
        anns_field="embedding",
        param=search_params,
        limit=top_k,
        output_fields=["text"]
    )
    retrieved = [hit.entity.get("text") for hit in results[0] if hit.entity]
    print("\n🔍 检索到的参考文档：")
    for i, t in enumerate(retrieved, 1):
        print(f"[{i}] {t[:100]}{'...' if len(t) > 100 else ''}")

    # 3. 拼接上下文
    context = "\n\n".join(retrieved)
    prompt = f"""你是一个知识问答助手。请根据以下参考资料回答问题。
参考资料：
{context}

问题：{question}

请结合参考资料，用简洁、准确的中文回答。
"""

    # 4. 调用 LLaMA2
    payload = {
        "model": llama2_model,
        "prompt": prompt,
        "stream": False
    }
    resp = requests.post(llama2_url, json=payload)
    resp.raise_for_status()
    answer = resp.json().get("response", "").strip()

    return answer

if __name__ == "__main__":
    connect_milvus()
    while True:
        q = input("\n请输入你的问题（输入 exit 退出）：").strip()
        if q.lower() in {"exit", "quit"}:
            break
        ans = rag_query(q)
        print("\n💡 回答：", ans)
